package flowdroid.algo;

import pojo.MethodMarks;
import soot.*;
import soot.jimple.InvokeExpr;
import soot.jimple.JimpleBody;
import tuple.data.DataKind;
import tuple.data.NewDataRule;
import tuple.entity.EntityRule;
import tuple.entity.MyEntity;
//import tuple.permission.MarkExtractionRule;
import tuple.permission.PermissionRule;
import util.FlowdroidUtils;

import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import java.io.IOException;
import java.util.Arrays;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.stream.Collectors;

public class ReverseDFSImpl implements ReverseDFS {

    /**
     * visitedClazz: 某条数据流遍历过的 sootClass
     * visitedMethods: 某条数据流遍历过的 sootMethod
     */
    public final HashSet<SootClass> visitedClazz;
    public final HashSet<SootMethod> visitedMethods;
//    public static Map<SootMethod, Boolean> reverseVisited = new HashMap<>();

    public ReverseDFSImpl() throws IOException {
        this.visitedClazz = new HashSet<>();
        this.visitedMethods = new HashSet<>();
        NewDataRule.ensureDataLoaded();
        EntityRule.ensureEntityLoaded();
        PermissionRule.ensureLoaded();
//        MarkExtractionRule.ensureExtractionPtnLoaded();
    }

    private static final int maxStackSizeForClass = 3;
    private static final int maxStackSizeForClassField = 2;
    private static final int maxStackSizeForClassConstructor = 3;
    private static final int maxStackSizeForMethods = 5;

    // 1. 以【类】的粒度检查: 包名类名, 字段, 构造器和静态构造器(注意构造器和静态构造器都依赖于 checkMethodBody,不检查参数
    @Override
    public void checkSootClass(SootClass sootClass, MethodMarks methodMarks, int currentLevel) {
        if (visitedClazz.contains(sootClass) || currentLevel >= maxStackSizeForClass)
            return;

        // 1  检查类名
        checkClassName(sootClass, methodMarks);

        // 2 检查构造器或者静态构造器
        checkClassConstructor(sootClass, methodMarks, currentLevel);

        // 3 检查类字段,其实是做递归; 每个字段有 declaring class,我们只检查这些declaring class的类名(通过递归层数控制)
        if (sootClass.getFieldCount() > 0) {
            checkClassField(sootClass, methodMarks, currentLevel);
        }

        visitedClazz.add(sootClass);
    }

    @Override
    public void checkMethod(SootMethod sootMethod, MethodMarks methodMarks, int recursiveLevel) {
        if (recursiveLevel >= maxStackSizeForMethods || visitedMethods.contains(sootMethod))
            return;

        //1. method name
        String name = sootMethod.getName();
        if (name.length() > 1 && !sootMethod.isConstructor() && !name.contains("dummy")) {
            checkMethodName(sootMethod, methodMarks);
        }

        //2. body
        JimpleBody body = null;
        Set<SootMethod> calleeMethods = new HashSet<>();
        try {
            FlowdroidUtils.forceLoadBody(sootMethod);
            body = (JimpleBody) sootMethod.retrieveActiveBody();
        } catch (Exception ignored) {
        }
        if (body != null) {
            Set<SootMethod> invokedMethods = checkMethodBody(body, methodMarks);
            calleeMethods.addAll(invokedMethods);
        }

        //3. params
        if (sootMethod.getParameterCount() > 0) {
            checkParams(sootMethod, methodMarks, recursiveLevel);
        }

        //4. reverseCheckInvokedMethods
        for (SootMethod callee : calleeMethods) {
            checkMethod(callee, methodMarks, recursiveLevel + 1);
        }

        visitedMethods.add(sootMethod);
    }


    private void checkClassName(SootClass sootClass, MethodMarks methodMarks) {
        String className = sootClass.getName();
        Set<String> splitName = Arrays.stream(className.split("\\."))
                .filter(x -> x.length() > 1)
                .collect(Collectors.toSet());
        if (!splitName.isEmpty()) {
            Set<MyEntity> entities = EntityRule.checkClassNameAsEntity(splitName);
            methodMarks.entitySet.addAll(entities);

            Set<DataKind> dataKinds = NewDataRule.checkClassNameAsData(splitName);
            methodMarks.dataKinds.addAll(dataKinds);

            //condition 从类名上看不出来的
        }
    }


    private void checkClassConstructor(SootClass sootClass, MethodMarks methodMarks, int recursiveLevel) {
        if (recursiveLevel >= maxStackSizeForClassConstructor) //最多递归 maxStackSizeForConstructor 层(尽量少,不准确)
            return;
        Set<SootMethod> methods = sootClass.getMethods()
                .stream()
                .filter(m -> m.isConstructor() && !visitedMethods.contains(m))
                .collect(Collectors.toSet());
        for (SootMethod method : methods) {
            checkMethod(method, methodMarks, recursiveLevel + 2); //这里复用 checkMethod,但是构造器不应该和方法主体一样重要,所以没有+1，而是+2，早点从递归中脱离;
        }
    }

    // 这是一个递归方法, checkClassField
    private void checkClassField(SootClass sootClass, MethodMarks methodMarks, int recursiveLevel) {
        if (recursiveLevel >= maxStackSizeForClassField) //最多递归 MaxStackSizeForClazz 层(尽量少,不准确)
            return;
        Set<SootClass> filteredFieldClazz = sootClass.getFields().stream()
                .map(SootField::getDeclaringClass)
                .filter(declaringClass -> !visitedClazz.contains(declaringClass) && !declaringClass.equals(sootClass))
                .collect(Collectors.toSet());
        for (SootClass cl : filteredFieldClazz) {
            checkSootClass(cl, methodMarks, recursiveLevel + 1);
        }
    }


    private void checkMethodName(SootMethod sootMethod, MethodMarks methodMarks) {
        String name = sootMethod.getName();
        Set<MyEntity> entities = EntityRule.checkMethodNameAsEntity(name);
        methodMarks.entitySet.addAll(entities);

        Set<DataKind> dataKinds = NewDataRule.checkMethodNameAsData(name);
        methodMarks.dataKinds.addAll(dataKinds);

        Set<String> stringMarksSet = PermissionRule.checkMethodPermission(sootMethod);
//        Set<String> stringMarksSet = MarkExtractionRule.checkMethodMarks(sootMethod);
        if (stringMarksSet.isEmpty()) {
            if (dataKinds.contains(DataKind.ANDROID_PERMISSION)) {
                stringMarksSet.add("android.permission");
            } else if (dataKinds.contains(DataKind.PERMISSION)) {
                stringMarksSet.add("permission");
            }
        }
        methodMarks.stringMarks.addAll(stringMarksSet);
    }

    //todo // 增加对 JAssignment JInstanceFieldRef 的处理
    private Set<SootMethod> checkMethodBody(@Nullable JimpleBody body, @Nonnull MethodMarks methodMarks) {
        if (body == null)
            return new HashSet<>();

        Set<Value> noDupValues = body.getUnits().stream()
                .flatMap(x -> x.getUseAndDefBoxes().stream())
                .map(x -> x.getValue())
                .collect(Collectors.toSet());//JimpleLocal, JSpecialInvokeExpre,JNewExpr,JInstanceFieldRef,ThisRef
//        Set<InvokeExpr> allInvokes = noDupValues
//                .stream()
//                .filter(x -> x instanceof InvokeExpr)
//                .map(x -> ((InvokeExpr) (x)))
//                .collect(Collectors.toSet());

        Set<MyEntity> entities = EntityRule.checkMethodBodyAsEntity(noDupValues);
        Set<DataKind> dataKinds = NewDataRule.checkMethodBodyAsData(noDupValues);// JimpleLocal,NULLCONstantt，intConstant,;StraingConstant;

        methodMarks.entitySet.addAll(entities);
        methodMarks.dataKinds.addAll(dataKinds);
        methodMarks.stringMarks = PermissionRule.checkMethodStmtsHasPermission(noDupValues);
//        methodMarks.stringMarks = MarkExtractionRule.checkMethodStmtsHasMarks(noDupValues);

        Set<SootMethod> invokedMethods = noDupValues.stream()
                .filter(x -> x instanceof InvokeExpr)
                .map(x -> ((InvokeExpr) (x)).getMethod())
                .filter(method -> !visitedMethods.contains(method)
                        && !method.getSignature().contains("java.lang")
                        && !method.getSignature().contains("java.util")
                )
                .collect(Collectors.toSet());
//
        return invokedMethods;
    }


    //todo
    private void checkParams(SootMethod sootMethod, MethodMarks methodMarks, int recursiveLevel) {
        List<Type> parameterTypes = sootMethod.getParameterTypes();
        // RefType
        for (Type parameterType : parameterTypes) {
            if (parameterType instanceof RefType) {
                RefType refType = (RefType) parameterType;
                SootClass sootClass = refType.getSootClass();
                if (!visitedClazz.contains(sootClass)) {
                    checkSootClass(sootClass, methodMarks, recursiveLevel + 1);
                }
            }
        }
    }
}
